"""
python lstm_test.py
"""

import numpy as np
import matplotlib.pyplot as plt
from keras.models import Sequential
from keras.layers import LSTM, Dense

# 生成正弦波加噪声的测试数据
np.random.seed(42)
time = np.arange(0, 100, 0.1)
sin_wave = np.sin(time)
noise = 0.1 * np.random.randn(len(time))
data = sin_wave + noise

# 准备数据
look_back = 10
X = []
y = []
for i in range(len(data) - look_back):
    X.append(data[i:i + look_back])
    y.append(data[i + look_back])
X = np.array(X)
y = np.array(y)

# 重塑输入数据以适应 LSTM 输入要求
X = np.reshape(X, (X.shape[0], X.shape[1], 1))

# 创建 LSTM 模型
model = Sequential()
model.add(LSTM(50, input_shape=(look_back, 1)))
model.add(Dense(1))

# 编译模型
model.compile(loss='mean_squared_error', optimizer='adam')

# 拟合模型
model.fit(X, y, epochs=50, batch_size=32, verbose=0)

# 进行预测
predictions = []
last_sequence = data[-look_back:]
for _ in range(50):
    new_prediction = model.predict(last_sequence.reshape(1, look_back, 1))[0][0]
    predictions.append(new_prediction)
    last_sequence = np.append(last_sequence[1:], new_prediction)

# 扩展时间序列以匹配预测数据长度
extended_time = np.arange(time[-1]+0.1, time[-1]+0.1+len(predictions)*0.1, 0.1)

# 绘制原始数据和预测数据的图像
plt.figure(figsize=(12, 6))
plt.plot(time[:len(data)], data, label='Original Data')
plt.plot(extended_time, predictions, label='Predicted Data')
plt.legend()
plt.xlabel('Time')
plt.ylabel('Value')
plt.title('LSTM Prediction on Sinusoidal Wave with Noise')
plt.show()
